﻿/**
 * @file zmq_service.hpp
 * @author Yunfeng.Bo (Yunfeng.Bo@qq.com)
 * @brief 基于 ZeroMQ 的服务，各模块可继承该服务，将模块功能通过异步的消息队列服务提供
 * @version 0.1
 * @date 2022-03-23
 *
 * 序列化使用 msgpack：
 *   - 不使用 boost.serializer 或自定义的序列化等 C++
 * 特有的序列化库，是为了兼容更多语言；
 *   - 不使用基于 json 的序列化库，是考虑性能问题；
 *   - 不使用 protobuf 等，是因为考虑到 msgpack 开发更简单一点，不需要 proto
 * 等接口定义文件；
 *
 * @copyright Copyright (c) 2022
 *
 */

#pragma once

#ifndef USE_PRECOMPILE_HEADER
#    include "pch.hpp"
#endif

#include "http_rep.h"
#include "serializer.hpp"

#ifdef USE_LIB_ZMQ
#    include <zmq.hpp>
#endif

// ZeroMQ 命令返回 response code
enum class response_code {
    SUCCESS              = 0,       // 调用成功
    FUNCTION_NOT_REGUIST = -1000,   // 调用的方法未注册
    SEND_ERROR           = -1001,   // 发送错误
    RECV_TIMEOUT         = -1002,   // 调用超时
    ARGS_COUNT_ERR = -1003,   // 参数不足（包括没指定方法名，也认为“参数不足”）
    ARGS_TYPE_ERR = -1004,   // 参数类型错误，无法正确转换为注册方法的参数类型
    RET_TYPE_ERR  = 1005,    // 返回值类型错误
    UNKNOWN       = -1999,   // 发生未知类型错误
};

#ifdef USE_LIB_ZMQ
/**
 * @brief ZeroMQ 命令调用客户端
 */
class zmq_req_client {
public:
    explicit zmq_req_client(zmq::context_t& ctx, const std::string& addr, int timeout = 5000) {
        try {
            m_socket = std::unique_ptr<zmq::socket_t, std::function<void(zmq::socket_t*)>>(
                new zmq::socket_t(ctx, zmq::socket_type::req), [](zmq::socket_t* sock) {
                    sock->close();
                    delete sock;
                    sock = nullptr;
                });
            m_socket->connect(addr);
            m_socket->set(zmq::sockopt::rcvtimeo, timeout);
        } catch (zmq::error_t& e) {
            LOG_ERROR << "Init ZeroMQ REQ client failed, error code: " << e.num()
                      << ", message: " << e.what();
        } catch (...) {
            LOG_ERROR << "Init ZeroMQ REQ client failed, unknown error;";
        }
    }
    virtual ~zmq_req_client() { m_socket->close(); }

    /**
     * @brief 调用 ZeroMQ 服务
     * @tparam R 返回值类型
     * @tparam ...Args
     * @param name RPC 方法名
     * @param ...args
     * @return response_t 类型，包含返回的 code、message 和 value
     */
    template<typename R, typename... Args>
    response_t<R> call(Args... args) {
        // 参数打包成 std::tuple
        using args_type = std::tuple<typename std::decay<Args>::type...>;
        args_type as    = std::make_tuple(args...);

        // 将 tuple 序列化成二进制数据
        msgpack::serializer sr;
        package_args(sr, as);

        // 将二进制化的参数放到 zmq 消息中
        zmq::message_t request(sr.size() + 1);
        memcpy(request.data(), sr.data(), sr.size());

        // 准备返回用数据结构
        response_t<R>  response;
        zmq::message_t reply;

        try {
            // 发送命令
            zmq::send_result_t ret_send = m_socket->send(request, zmq::send_flags::none);
            if (!ret_send.has_value()) {
                LOG_ERROR << "zmq_req_client call failed, can not send data.";
                // 返回错误信息
                response.code((int)response_code::SEND_ERROR);
                response.message("zmq_req_client call failed, can not send data.");
                return response;
            }

            // 接收返回数据
            zmq::recv_result_t ret_recv = m_socket->recv(reply, zmq::recv_flags::none);
            if (!ret_recv.has_value()) {
                // 接收数据超时，返回错误信息
                response.code((int)response_code::RECV_TIMEOUT);
                response.message("recive reply timeout");
                return response;
            }
        } catch (zmq::error_t& e) {
            // 上下文context已被中止
            LOG_ERROR << "zmq_req_client proecess error, code: " << e.num()
                      << ", message: " << e.what();
        } catch (...) {
            // 未知错误，有可能是 onrequest 时发生，且没有在处理方法内正确处理
            LOG_ERROR << "zmq_req_client proecess error.";
        }

        // 解析到 response
        msgpack::serializer res((const char*)reply.data(), reply.size());
        res >> response;

        return response;
    }

private:
    // 用于发送命令的接口
    std::unique_ptr<zmq::socket_t, std::function<void(zmq::socket_t*)>> m_socket;
};

/**
 * @brief ZeroMQ 消息订阅客户端
 */
class zmq_sub_client {
public:
    explicit zmq_sub_client(zmq::context_t& ctx, const std::string& addr) {
        try {
            m_socket = std::unique_ptr<zmq::socket_t, std::function<void(zmq::socket_t*)>>(
                new zmq::socket_t(ctx, zmq::socket_type::sub), [](zmq::socket_t* sock) {
                    sock->close();
                    delete sock;
                    sock = nullptr;
                });
            m_socket->connect(addr);
        } catch (zmq::error_t& e) {
            LOG_ERROR << "Init ZeroMQ SUB client failed, error code: " << e.num()
                      << ", message: " << e.what();
        } catch (...) {
            LOG_ERROR << "Init ZeroMQ SUB client failed, unknown error;";
        }
    }
    virtual ~zmq_sub_client() { m_socket->close(); }

    /**
     * @brief 设置需要订阅的“标题”，可以多次调用这个方法，设置多个需要订阅的标题
     * @param envelope 需要订阅的“标题”，如果设置为空字符串（或者不设置），默认接收所有消息
     */
    void subscribe(const std::string& envelope) {
        m_socket->set(zmq::sockopt::subscribe, envelope);
    }

    //// Receive all parts of the message
    // std::vector<zmq::message_t> recv_msgs;
    // zmq::recv_result_t result = zmq::recv_multipart(subscriber, std::back_inserter(recv_msgs));
    // assert(result && "recv failed");
    // assert(*result == 2);

    // std::cout << "Thread2: [" << recv_msgs[0].to_string() << "] " << recv_msgs[1].to_string()
    //          << std::endl;

private:
    // 用于发送命令的接口
    std::unique_ptr<zmq::socket_t, std::function<void(zmq::socket_t*)>> m_socket;
};
#endif   // USE_LIB_ZMQ

/**
 * @brief 内部管理用，注册为 可被远程调用（RPC）的 服务函数
 *          第一个参数用于返回调用结果
 *          第二个参数是调用的参数
 */
typedef std::function<void(msgpack::serializer*, msgpack::serializer&)> rpc_func;

template<typename T>
struct type_xx {
    typedef T type;
};

template<>
struct type_xx<void> {
    typedef int8_t type;
};

// 用 Tuple 做参数调用函数模板类，辅助
template<typename Function, typename Tuple, std::size_t... Index>
decltype(auto) invoke_impl(Function&& func, Tuple&& t, std::index_sequence<Index...>) {
    return func(std::get<Index>(std::forward<Tuple>(t))...);
}

// 用 Tuple 做参数调用函数模板类
template<typename Function, typename Tuple>
decltype(auto) invoke(Function&& func, Tuple&& t) {
    constexpr auto size = std::tuple_size<typename std::decay<Tuple>::type>::value;
    return invoke_impl(
        std::forward<Function>(func), std::forward<Tuple>(t), std::make_index_sequence<size>{});
}

// 调用服务方法
template<typename R, typename F, typename ArgsTuple>
typename std::enable_if<!std::is_same<R, void>::value, typename type_xx<R>::type>::type
call_helper(F f, ArgsTuple args) {
    return invoke(f, args);
}

/**
 * @brief 基于 ZeroMQ 的服务基类
 */
class zmq_service {
public:
    explicit zmq_service(const std::string& service_name)
        : m_service_name(service_name)
#ifdef USE_LIB_ZMQ
        , m_ctx(1)
        , m_rep_socket(nullptr)
        , m_pub_socket(nullptr)
#endif   // USE_LIB_ZMQ
        , m_stop_running(true)
        , m_work_thread(nullptr) {

#ifdef USE_LIB_ZMQ
        m_ctx.set(zmq::ctxopt::blocky, false);
#endif   // USE_LIB_ZMQ
    }
    virtual ~zmq_service() {
#ifdef USE_LIB_ZMQ
        stop_serivce();
#endif   // USE_LIB_ZMQ
    }

    /**
     * @brief 通过直接函数调用的方式，使用模块功能
     */
    msgpack::serializer::ptr call(msgpack::serializer& sr) {
        msgpack::serializer sr_request(sr.data(), sr.size());

        // 先取出调用的方法名称
        std::string funcname;

        // 封装返回的数据
        msgpack::serializer::ptr rptr(nullptr);

        try {
            sr_request >> funcname;

            // 尝试调用名称指定的服务方法，并获得返回的 response
            rptr = call_(funcname, sr_request);
        } catch (error_args_count& er) {
            // 参数错误（或者甚至没有方法名）
            LOG_WARNING << "Call args not enough, func: " << funcname
                        << ", error mesgae: " << er.what();
            rptr = std::make_shared<msgpack::serializer>();
            *rptr.get() << static_cast<int>(response_code::ARGS_COUNT_ERR);
            *rptr.get() << std::string(er.what());
        } catch (error_args_type& er) {
            // 参数类型错误（无法转换）
            LOG_WARNING << "Call args type error, func: " << funcname
                        << ", error mesgae: " << er.what();
            rptr = std::make_shared<msgpack::serializer>();
            *rptr.get() << static_cast<int>(response_code::ARGS_TYPE_ERR);
            *rptr.get() << std::string(er.what());
        } catch (...) {
            // 其它未知错误
            LOG_WARNING << "Call unknown error, func: " << funcname;
            rptr = std::make_shared<msgpack::serializer>();
            *rptr.get() << static_cast<int>(response_code::UNKNOWN);
            *rptr.get() << std::string("Unknown error!");
        }

        return rptr;
    }

    /**
     * @brief 注册一个普通函数成为 RPC 服务方法
     * @tparam F
     * @param name client 调用时使用的方法名
     * @param func 服务响应方法
     */
    template<typename F>
    void regist(const std::string& name, F func) {
        m_map_funcs[name] = std::bind(
            &zmq_service::callproxy<F>, this, func, std::placeholders::_1, std::placeholders::_2);
    }

    /**
     * @brief 注册一个类成员函数成为 RPC 服务方法
     * @tparam F
     * @tparam S
     * @param name client 调用时使用的方法名
     * @param func 服务响应方法，类成员函数 &class::function
     * @param s 提供响应的类实例 &class_object
     */
    template<typename F, typename S>
    void regist(const std::string& name, F func, S* s) {
        m_map_funcs[name] = std::bind(&zmq_service::callproxy<F, S>,
                                      this,
                                      func,
                                      s,
                                      std::placeholders::_1,
                                      std::placeholders::_2);
    }

    /**
     * @brief 与 regist 功能相对，注销一个远程方法
     * @param name client 调用时使用的方法名
     */
    void unregist(const std::string& name) {
        m_map_funcs.erase(name);
    }

#ifdef USE_LIB_ZMQ
    /**
     * @brief 通过 PUB socket 广播发送一条消息数据
     * @param envelope 消息标题，用于消息订阅者过滤需要的消息
     * @param message 待发送的消息
     */
    virtual int pub_message(const std::string& envelope, zmq::message_t& message) {
        if (!m_pub_socket) {
            LOG_ERROR << "PUB Socket not inited!";
            return -1;
        }

        try {
            m_pub_socket->send(zmq::buffer(envelope), zmq::send_flags::sndmore);
            zmq::send_result_t ret = m_pub_socket->send(message, zmq::send_flags::none);
            if (!ret.has_value()) {
                // 没发送成功，通常是超时什么的
                LOG_ERROR << "PUB send message failed!";
                return -1;
            }
        } catch (zmq::error_t& e) {
            LOG_ERROR << "PUB send message failed! error code: " << e.num()
                      << ", message: " << e.what();
        } catch (...) {
            LOG_ERROR << "PUB send message failed, unknown error.";
        }

        return 0;
    }

    /**
     * @brief 处理 zmq_service 接收到的请求
     *          具体的服务功能模块，可以重载这个方法，实现在低层响应 req/rep
     *          例如一个 http server 模块，就可以在低层直接转发 request 和 reply 的方式实现命令转发
     *          如果不重载这个方法，则可以通过注册的方式实现响应方法的调用
     * @param request in, ZeroMQ 消息结构的请求数据
     * @param reply out, ZeroMQ 消息结构的返回数据
     */
    virtual void on_request(const zmq::message_t& request, zmq::message_t& reply) {
        // RPC 方法名
        std::string funcname;

        // 封装返回的数据
        msgpack::serializer::ptr rptr(nullptr);

        try {
            // 准备序列化
            msgpack::serializer sr((char*)request.data(), request.size());

            // 先取出调用的方法名称
            sr >> funcname;

            // 尝试调用名称指定的服务方法，并获得返回的 response
            rptr = call_(funcname, sr);

            LOG_INFO << "ZMQ_RPC Call finished, func: " << funcname;
        } catch (error_args_count er) {
            // 参数错误（或者甚至没有方法名）
            LOG_WARNING << "ZMQ_RPC Call args not enough, func: " << funcname
                        << ", error mesgae: " << er.what();
            rptr = std::make_shared<msgpack::serializer>();
            *rptr.get() << static_cast<int>(response_code::ARGS_COUNT_ERR);
            *rptr.get() << std::string(er.what());
        } catch (error_args_type er) {
            // 参数类型错误（无法转换）
            LOG_WARNING << "ZMQ_RPC Call args type error, func: " << funcname
                        << ", error mesgae: " << er.what();
            rptr = std::make_shared<msgpack::serializer>();
            *rptr.get() << static_cast<int>(response_code::ARGS_TYPE_ERR);
            *rptr.get() << std::string(er.what());
        } catch (...) {
            // 其它未知错误
            LOG_WARNING << "ZMQ_RPC Call unknown error, func: " << funcname;
            rptr = std::make_shared<msgpack::serializer>();
            *rptr.get() << static_cast<int>(response_code::UNKNOWN);
            *rptr.get() << std::string("Unknown error!");
        }

        // 发送 response，如果返回数据为空，则不返回，客户端需要正确处理超时
        if (rptr) {
            reply.rebuild(rptr->size());
            memcpy((void*)reply.data(), rptr->data(), rptr->size());
        } else {
            LOG_ERROR << "ZMQ_RPC Call, but somthing error, no response!";
        }
    }

    /**
     * @brief 使用指定的 ZeroMQ context 启动服务
     * @param ctx ZeroMQ context
     * @param REP_addr 请求服务地址
     * @param PUB_addr 消息通知地址，如果为空，则不会创建消息接口，只有命令接口
     * @return 创建成功返回 0
     */
    int start_service(zmq::context_t& ctx, const std::string& REP_addr,
                      const std::string& PUB_addr) {
        if (m_work_thread) {
            LOG_ERROR << "ZeroMQ Service [" << m_service_name
                      << "] is running, if need restart, stop service first!";
            return 0;
        }

        int ret(0);
        // 初始化命令 socket
        ret = init_rep_socket(ctx, REP_addr);
        if (ret != 0)
            return ret;

        // 初始化消息 socket
        if (!PUB_addr.empty()) {
            ret = init_pub_socket(ctx, PUB_addr);
            if (ret != 0)
                return ret;
        }

        m_stop_running = false;
        m_work_thread  = std::make_shared<std::thread>([&]() {
            // 设置超时 100 ms，这样每隔一段时间就有机会判断一下是不是该退出了
            m_rep_socket->set(zmq::sockopt::rcvtimeo, 100);
            while (!m_stop_running) {
                try {
                    // 接收命令数据
                    zmq::message_t     request;
                    zmq::recv_result_t ret_recv =
                        m_rep_socket->recv(request, zmq::recv_flags::none);
                    if (!ret_recv.has_value()) {
                        // 通常是超时了
                        continue;
                    }

                    // 响应命令
                    zmq::message_t reply;
                    on_request(request, reply);

                    // 返回
                    m_rep_socket->send(reply, zmq::send_flags::dontwait);
                } catch (zmq::error_t& e) {
                    LOG_ERROR << "ZeroMQ Service [" << m_service_name
                              << "] proecess error, code: " << e.num() << ", message: " << e.what();

                    // 上下文context已被中止
                    if (e.num() == ETERM)
                        break;
                } catch (...) {
                    // 未知错误，有可能是 onrequest 时发生，且没有在处理方法内正确处理
                    LOG_ERROR << "ZeroMQ Service [" << m_service_name << "] proecess unknown error";
                }
            }
        });

        LOG_INFO << "ZeroMQ Service [" << m_service_name << "] start!";

        return 0;
    }

    /**
     * @brief 使用内部 zeromq context 启动服务（不能使用 inproc）
     * @param REP_addr 请求服务地址（命令）
     * @param PUB_addr 消息发布地址（消息）
     * @return 创建成功返回 0
     */
    int start_service(const std::string& REQ_addr, const std::string& PUB_addr) {
        return start_service(m_ctx, REQ_addr, PUB_addr);
    }

    /**
     * @brief 停止服务
     * @return
     */
    int stop_serivce() {
        // 停止线程工作
        if (m_work_thread) {
            m_stop_running = true;
            if (m_work_thread->joinable())
                m_work_thread->join();
            m_work_thread = nullptr;
        }

        // 关闭响应命令接口
        if (m_rep_socket) {
            m_rep_socket->close();
            m_rep_socket = nullptr;
        }

        // 关闭消息发送接口
        if (m_pub_socket) {
            m_pub_socket->close();
            m_pub_socket = nullptr;
        }

        // 销毁上下文
        m_ctx.close();

        LOG_INFO << "ZeroMQ Service [" << m_service_name << "] stoped!";

        return 0;
    }

private:
    int init_rep_socket(zmq::context_t& ctx, const std::string& REP_addr) {
        try {
            if (m_rep_socket)
                m_rep_socket->close();

            m_rep_socket = std::unique_ptr<zmq::socket_t, std::function<void(zmq::socket_t*)>>(
                new zmq::socket_t(ctx, zmq::socket_type::rep), [](zmq::socket_t* sock) {
                    sock->close();
                    delete sock;
                    sock = nullptr;
                });
            m_rep_socket->bind(REP_addr);
            LOG_INFO << "ZeroMQ stocker bind at : " << REP_addr;
        } catch (zmq::error_t& e) {
            LOG_ERROR << "ZeroMQ stocker bind at : " << REP_addr
                      << " failed, error code: " << e.num() << ", error message: " << e.what();
            return -1;
        } catch (...) {
            LOG_ERROR << "ZeroMQ stocker bind at : " << REP_addr << " failed!";
            return -1;
        }

        return 0;
    }

    int init_pub_socket(zmq::context_t& ctx, const std::string& PUB_addr) {
        try {
            if (m_pub_socket)
                m_pub_socket->close();

            m_pub_socket = std::unique_ptr<zmq::socket_t, std::function<void(zmq::socket_t*)>>(
                new zmq::socket_t(ctx, zmq::socket_type::pub), [](zmq::socket_t* sock) {
                    sock->close();
                    delete sock;
                    sock = nullptr;
                });
            m_pub_socket->bind(PUB_addr);
            LOG_INFO << "ZeroMQ stocker bind at : " << PUB_addr;
        } catch (zmq::error_t& e) {
            LOG_ERROR << "ZeroMQ stocker bind at : " << PUB_addr
                      << " failed, error code: " << e.num() << ", error message: " << e.what();
            return -1;
        } catch (...) {
            LOG_ERROR << "ZeroMQ stocker bind at : " << PUB_addr << " failed!";
            return -1;
        }

        return 0;
    }
#endif   // USE_LIB_ZMQ

private:
    /**
     * @brief server 端，根据方法名称，查找并调用注册的服务方法
     * @param name 被调用的RPC方法名
     * @param sr_args 序列化的参数
     * @return 返回打包好的（序列化后）的 response
     */
    msgpack::serializer::ptr call_(const std::string& name, msgpack::serializer& sr_args) {
        msgpack::serializer::ptr sp = std::make_shared<msgpack::serializer>();

        if (m_map_funcs.find(name) == m_map_funcs.end()) {
            LOG_WARNING << "RPC call, but function not regist, name = " << name << std::endl;
            *sp.get() << static_cast<int>(response_code::FUNCTION_NOT_REGUIST);
            *sp.get() << std::string("function not regist : " + name);
            return sp;
        }

        auto func = m_map_funcs[name];

        // 调用服务方法
        func(sp.get(), sr_args);

        return sp;
    }

    template<typename F>
    void callproxy(F func, msgpack::serializer* pr, msgpack::serializer& sr_args) {
        callproxy_(func, pr, sr_args);
    }

    template<typename F, typename S>
    void callproxy(F func, S* s, msgpack::serializer* pr, msgpack::serializer& sr_args) {
        callproxy_(func, s, pr, sr_args);
    }

    template<typename R, typename... Args>
    void callproxy_(R (*func)(Args...), msgpack::serializer* pr, msgpack::serializer& sr_args) {
        callproxy_(std::function<R(Args...)>(func), pr, sr_args);
    }

    template<typename R, typename C, typename S, typename... Args>
    void callproxy_(R (C::*func)(Args...), S* s, msgpack::serializer* pr,
                    msgpack::serializer& sr_args) {
        using args_type = std::tuple<typename std::decay<Args>::type...>;

        args_type as = sr_args.get_tuple<args_type>(std::index_sequence_for<Args...>{});

        auto                      ff = [=](Args... args) -> R { return (s->*func)(args...); };
        typename type_xx<R>::type r  = call_helper<R>(ff, as);

        response_t<R> response((int)response_code::SUCCESS, "success", r);
        (*pr) << response;
    }

    template<typename R, typename... Args>
    void callproxy_(std::function<R(Args... args)> func, msgpack::serializer* pr,
                    msgpack::serializer& sr_args) {
        using args_type = std::tuple<typename std::decay<Args>::type...>;

        args_type as = sr_args.get_tuple<args_type>(std::index_sequence_for<Args...>{});

        typename type_xx<R>::type r = call_helper<R>(func, as);

        response_t<R> response((int)response_code::SUCCESS, "success", r);
        (*pr) << response;
    }

private:
    // 服务名称，只是记录日志什么的方便
    std::string m_service_name;

    // 方法名 => 方法实现 map
    std::map<std::string, rpc_func> m_map_funcs;

#ifdef USE_LIB_ZMQ
    // 在不指定外部context时，使用这个内部的context
    zmq::context_t m_ctx;

    // 用于响应命令的接口
    std::unique_ptr<zmq::socket_t, std::function<void(zmq::socket_t*)>> m_rep_socket;

    // 用于响应命令的接口
    std::unique_ptr<zmq::socket_t, std::function<void(zmq::socket_t*)>> m_pub_socket;
#endif   // USE_LIB_ZMQ

    // 服务线程停止标志
    bool m_stop_running;

    // 服务线程
    std::shared_ptr<std::thread> m_work_thread;
};